# -*- coding: utf-8 -*-
"""
# @FileName:     image_utils.py
# @AuthorName:   Sanqi Lu (Lingwei Dang)
# @Institution:  SCUT, Guangzhou, China
# @EmailAddress: lenvondang@163.com
# @CreateTime:   2024/12/21 13:47
"""

import numpy as np
import cv2


def draw_hoi_pose_on_given_img_rendered_by_depth(hoi_pose, given_img_rgb):
    '''
    Args:
        hand_pose: [2, 21+128, 3]
        given_img_rgb: [H, W, 3]

    '''
    img_wh = given_img_rgb.shape[1], given_img_rgb.shape[0]
    framebgr = cv2.cvtColor(given_img_rgb, cv2.COLOR_RGB2BGR)
    left_hand_color = (0, 255, 255) # (0, 0, 255)
    right_hand_color = (0, 255, 255) # (0, 0, 255)
    left_line_color = (0, 0, 255)
    right_line_color = (0, 255, 0) # (0, 255, 255)
    tool_color = (218, 53, 237) 
    target_color = (224, 38, 21)
    Is = [0, 1, 2, 3, 0, 5, 6, 7, 9, 10, 11, 13, 14, 15, 0, 17, 18, 19, 5, 9, 13]
    Js = [1, 2, 3, 4, 5, 6, 7, 8, 10, 11, 12, 14, 15, 16, 17, 18, 19, 20, 9, 13, 17]
    
    hand_pose = hoi_pose[:, :21, :]
    obj_pose = hoi_pose[:, 21:, :]
    all_elems = []
    # 处理手部Pose
    # hand_pose = np.around(hand_pose).astype(int)
    xs = hand_pose[:, :, 0].round().astype(int)
    ys = hand_pose[:, :, 1].round().astype(int)
    zs = hand_pose[:, :, 2]
    xs[xs<0] = 0
    xs[xs>=img_wh[0]] = img_wh[0] - 1
    ys[ys < 0] = 0
    ys[ys >= img_wh[1]] = img_wh[1] - 1
    # 深度排序
    for hand_idx in range(2):
        for joint_idx in range(hand_pose.shape[1]):
            all_elems.append({
                "elem_type": "point",
                "hoi_type": "right_hand" if hand_idx == 0 else "left_hand",
                "color": right_hand_color if hand_idx == 0 else left_hand_color,
                "depth": zs[hand_idx, joint_idx],
                "point_xy": (xs[hand_idx, joint_idx], ys[hand_idx, joint_idx])
            })
        for line_idx in range(len(Is)):
            start_point = (xs[hand_idx, Is[line_idx]], ys[hand_idx, Is[line_idx]])
            end_point = (xs[hand_idx, Js[line_idx]], ys[hand_idx, Js[line_idx]])
            all_elems.append({
                "elem_type": "line",
                "hoi_type": "right_line" if hand_idx == 0 else "left_idx",
                "color": right_line_color if hand_idx == 0 else left_line_color,
                "depth": (zs[hand_idx, Is[line_idx]] + zs[hand_idx, Js[line_idx]]) / 2,
                "point_xy": [start_point, end_point]
            })
            
    # clip
    xs = obj_pose[:, :, 0].round().astype(int)
    ys = obj_pose[:, :, 1].round().astype(int)
    zs = obj_pose[:, :, 2]
    xs[xs<0] = 0
    xs[xs>=img_wh[0]] = img_wh[0] - 1
    ys[ys < 0] = 0
    ys[ys >= img_wh[1]] = img_wh[1] - 1
    for obj_idx in range(2):
        for joint_idx in range(obj_pose.shape[1]):
            all_elems.append({
                "elem_type": "point",
                "hoi_type": "tool" if obj_idx == 0 else "target",
                "color": tool_color if obj_idx == 0 else target_color,
                "depth": zs[obj_idx, joint_idx],
                "point_xy": (xs[obj_idx, joint_idx], ys[obj_idx, joint_idx])
            })
    
    # 按深度从远到近排序
    all_elems.sort(key=lambda e: -e["depth"])
    
    # 按深度顺序渲染元素
    for element in all_elems:
        if element["elem_type"] == "point":
            if "_hand" in element["hoi_type"]:
                cv2.circle(framebgr, element["point_xy"], radius=2, color=element["color"], thickness=1)
            else:
                cv2.circle(framebgr, element["point_xy"], radius=2, color=element["color"], thickness=2)
        elif element["elem_type"] == "line":
            cv2.line(framebgr, element["point_xy"][0], element["point_xy"][1], color=element["color"], thickness=2)
        # print(f'element["hoi_type"]: {element["hoi_type"]}, element["depth"]: {element["depth"]}')

    frame = cv2.cvtColor(framebgr, cv2.COLOR_BGR2RGB)
    return frame


    
def draw_hand_pose_on_given_img(hand_pose, given_img_rgb, handtype="right"):
    '''

    :param hand_pose: [21, 3]
    :param given_img_rgb: [h, w, 3]
    :param handtype:
    :return:
    '''
    assert hand_pose.shape[0] == 21
    Is = [0, 1, 2, 3, 0, 5, 6, 7, 9, 10, 11, 13, 14, 15, 0, 17, 18, 19, 5, 9, 13]
    Js = [1, 2, 3, 4, 5, 6, 7, 8, 10, 11, 12, 14, 15, 16, 17, 18, 19, 20, 9, 13, 17]
    img_wh = given_img_rgb.shape[1], given_img_rgb.shape[0]
    # print(f"hand_pose: {hand_pose.min()}, {hand_pose.max()}")
    # 四舍五入取整数
    hand_pose = np.around(hand_pose).astype(int)
    # clip
    xs = hand_pose[:, 0]
    ys = hand_pose[:, 1]
    # print(f"裁剪前")
    # print(f"xs: {xs}")
    # print(f"ys: {ys}")

    xs[xs<0] = 0
    xs[xs>=img_wh[0]] = img_wh[0] - 1
    ys[ys < 0] = 0
    ys[ys >= img_wh[1]] = img_wh[1] - 1
    # print(f"裁剪后")
    # print(f"xs: {xs}")
    # print(f"ys: {ys}")
    framebgr = cv2.cvtColor(given_img_rgb, cv2.COLOR_RGB2BGR)
    # print(framebgr.shape)
    # 绘制点坐标, color="r", s=120
    for pi in range(hand_pose.shape[0]):
        center = (xs[pi], ys[pi])
        cv2.circle(framebgr, center, radius=2, color=(0, 0, 255), thickness=2)
    # 画线：color="#000000"
    for idx in range(len(Is)):
        start_point = (xs[Is[idx]], ys[Is[idx]])
        end_point = (xs[Js[idx]], ys[Js[idx]])
        cv2.line(framebgr, start_point, end_point, color=(0, 255, 255) if handtype == "right" else (0, 255, 0),
                 thickness=2)

    frame = cv2.cvtColor(framebgr, cv2.COLOR_BGR2RGB)
    return frame

def draw_obj_pose_on_given_img(obj_pose, given_img_rgb, obj_type="tool"):
    '''

    :param hand_pose: [n, 3]
    :param given_img_rgb: [h, w, 3]
    :param handtype:
    :return:
    '''
    # print(f"min max: {obj_pose[:, 0].min()}, {obj_pose[:, 0].max()}, y: {obj_pose[:, 1].min()}, {obj_pose[:, 1].max()}")
    # #ed35da, #1526e0
    color_bgr = (218, 53, 237) if obj_type == "tool" else (224, 38, 21)
    assert obj_pose.shape[-1] == 3 and len(obj_pose.shape) == 2

    img_wh = given_img_rgb.shape[1], given_img_rgb.shape[0]
    # print(f"hand_pose: {hand_pose.min()}, {hand_pose.max()}")
    # 四舍五入取整数
    obj_pose = np.around(obj_pose).astype(int)
    # clip
    xs = obj_pose[:, 0]
    ys = obj_pose[:, 1]
    # print(f"裁剪前")
    # print(f"xs: {xs}")
    # print(f"ys: {ys}")

    xs[xs<0] = 0
    xs[xs>=img_wh[0]] = img_wh[0] - 1
    ys[ys < 0] = 0
    ys[ys >= img_wh[1]] = img_wh[1] - 1
    # print(f"裁剪后")
    # print(f"xs: {xs}")
    # print(f"ys: {ys}")
    framebgr = cv2.cvtColor(given_img_rgb, cv2.COLOR_RGB2BGR)
    # print(framebgr.shape)
    # 绘制点坐标, color="r", s=120
    for pi in range(obj_pose.shape[0]):
        center = (xs[pi], ys[pi])
        cv2.circle(framebgr, center, radius=2, color=color_bgr, thickness=2)
  
    frame = cv2.cvtColor(framebgr, cv2.COLOR_BGR2RGB)
    return frame
